package org.dromara.easyai.extensions.cuda;

import jcuda.jcublas.JCublas;
import jcuda.jcudnn.JCudnn;
import jcuda.runtime.JCuda;
import org.dromara.easyai.matrixTools.CudaMatrix;
import org.dromara.easyai.matrixTools.Matrix;
import org.dromara.jcudax.JCudax;

import static org.dromara.easyai.extensions.cuda.MatrixUtil.*;


public class CudaMatrixImpl implements CudaMatrix {

	static {
		System.out.println("EasyAI CUDA-12.0.0 extensions loaded.");
	}

	@Override
	public void init() throws Exception {
		// 提前初始化 CUDA 上下文
		JCuda.initialize();
		JCublas.cublasInit();
		JCudnn.initialize();
		JCudax.initialize();
		System.out.println("EasyAI CUDA-12.0.0 extensions init success.");
	}

	public void destroy() throws Exception {
		// 释放 CUDA 上下文
		JCublas.cublasShutdown();
	}

	@Override
	public void softMax(Matrix a) throws Exception {
		float[] h_A = toRowMajorArray(a);
		float[] h_C = CudaFp32Util.matrixSoftmax(h_A, a.getX(), a.getY());
		setRowMajorArray(a, h_C);
	}

	@Override
	public Matrix matrixSoftMaxPd(Matrix qkt, Matrix errorMatrix, float wordVectorDimension) throws Exception {
		float[] h_qkt = toRowMajorArray(qkt);
		float[] h_errorMatrix = toRowMajorArray(errorMatrix);
		float[] h_grMatrix = new float[qkt.getX() * qkt.getY()];
		CudaFp32Util.matrixSoftMaxPd(h_qkt, h_errorMatrix, h_grMatrix, qkt.getX(), qkt.getY(), wordVectorDimension);
		return rowMajorArrayToMatrix(h_grMatrix, qkt.getX(), qkt.getY());
	}

	@Override
	public Matrix mulMatrix(Matrix a, Matrix b) throws Exception {
//		float[] h_A = toColumnMajorArray(a);
//		float[] h_B = toColumnMajorArray(b);
//		float[] h_C = CudaFp32Util.matrixMulMatrix(h_A, h_B, a.getX(), a.getY(), b.getY());
//		return columnMajorArrayToMatrix(h_C, a.getX(), b.getY());
		float[] h_A = a.getCudaMatrix();
		float[] h_B = b.getCudaMatrix();
		float[] h_C = CudaFp32Util.matrixMulMatrix(h_A, h_B, a.getX(), a.getY(), b.getY());

		Matrix res = new Matrix(a.getX(), b.getY());
		res.setCudaMatrix(h_C, a.getX(), b.getY());
		return res;
	}

	@Override
	public void mathAdd(Matrix a, float scalar) throws Exception {
//		float[] h_A = toColumnMajorArray(a);
//		float[] h_C = CudaFp32Util.matrixAddScalar(h_A, scalar, a.getX(), a.getY());
//		setColumnMajorArray(a, h_C);

		float[] h_A = a.getCudaMatrix();
		float[] h_C = CudaFp32Util.matrixAddScalar(h_A, scalar, a.getX(), a.getY());
		a.setCudaMatrix(h_C, a.getX(), a.getY());
	}

	@Override
	public void mathSub(Matrix a, float scalar) throws Exception {
//		float[] h_A = toColumnMajorArray(a);
//		float[] h_C = CudaFp32Util.matrixAddScalar(h_A, -scalar, a.getX(), a.getY());
//		setColumnMajorArray(a, h_C);

		float[] h_A = a.getCudaMatrix();
		float[] h_C = CudaFp32Util.matrixAddScalar(h_A, -scalar, a.getX(), a.getY());
		a.setCudaMatrix(h_C, a.getX(), a.getY());
	}

	@Override
	public void mathMul(Matrix a, float scalar) throws Exception {
//		float[] h_A = toColumnMajorArray(a);
//		float[] h_C = CudaFp32Util.matrixScale(h_A, scalar, a.getX(), a.getY());
//		setColumnMajorArray(a, h_C);

		float[] h_A = a.getCudaMatrix();
		float[] h_C = CudaFp32Util.matrixScale(h_A, scalar, a.getX(), a.getY());
		a.setCudaMatrix(h_C, a.getX(), a.getY());
	}

	@Override
	public void mathDiv(Matrix a, float scalar) throws Exception {
//		float[] h_A = toColumnMajorArray(a);
//		float[] h_C = CudaFp32Util.matrixScale(h_A, 1/scalar, a.getX(), a.getY());
//		setColumnMajorArray(a, h_C);

		float[] h_A = a.getCudaMatrix();
		float[] h_C = CudaFp32Util.matrixScale(h_A, 1/scalar, a.getX(), a.getY());
		setColumnMajorArray(a, h_C);
	}
}
